/**
 * 
 */
package edu.umd.clip.lm.programs;

import java.io.*;
import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.lang.management.MemoryUsage;
import java.nio.channels.FileChannel;
import java.text.DateFormat;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Date;
import java.util.Timer;
import java.util.TimerTask;

import edu.berkeley.nlp.util.*;

import edu.umd.clip.jobs.JobManager;
import edu.umd.clip.lm.factors.FactorTuple;
import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.model.training.*;
import edu.umd.clip.lm.util.IO;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class LMNewSmootherEM {
	public static class Options {
        @Option(name = "-config", required = true, usage = "XML config file")
		public String config;

        @Option(name = "-jobs", required = false, usage = "number of concurrent jobs (default: 1)")
        public int jobs = 1;
        @Option(name = "-minstep", required = false, usage = "minimum change between iterations")
        public double minstep = 0.005;
        @Option(name = "-maxiter", required = false, usage = "maximum number of EM iterations")
        public int maxiter = 100;
        @Option(name = "-bucketSize", required = false, usage = "number of events per bucket (default: 0 (no bucketing))")
        public int bucketSize = 0;
        @Option(name = "-meminfo", required = false, usage = "track memory usage (default: false)")
        public boolean meminfo = false;
        @Option(name = "-lm", required = false, usage = "LM ID to train (default: " + LanguageModel.PRIMARY_LM_ID + ")")
		public String lm = LanguageModel.PRIMARY_LM_ID;        
        @Option(name = "-dbdir", required = true, usage = "the directory for the temporary database")
		public String dbdir;        
        @Option(name = "-datadir", required = true, usage = "the directory for the temporary files")
		public String datadir;        
        @Option(name = "-prune-tree", required = false, usage = "prune the decision tree (default: false)")
		public boolean pruneTree = false;        
        @Option(name = "-skip-data", required = false, usage = "skip data population step (default: false)")
		public boolean skipData = false;        
        @Option(name = "-singleFold", required = false, usage = "use a single fold (default: -1 (use cross-validation))")
        public int singleFold = -1;
        @Option(name = "-use-lambdas", required = false, usage = "initialize lamdas to values in the tree (default: false)")
        public boolean useLambdas = false;
        @Option(name = "-dummy-smoother", required = false, usage = "use unsmoothed dev data (default: false)")
        public boolean dummySmoother = false;
        @Option(name = "-dump-lambdas-prefix", required = false, usage = "the file prefix to dump lambdas (default: none)")
        public String dumpLambdasPrefix = null;
        @Option(name = "-discountScale", required = false, usage = "discount scale (default: 1.0)")
        public double discountScale = 1.0;
        @Option(name = "-dump-ngrams", required = false, usage = "dump ngram count (default: false)")
        public boolean dumpNgrams = false;
	}

	/**
	 * @param args
	 * @throws Exception 
	 */
	public static void main(String[] args) throws Exception {
        OptionParser optParser = new OptionParser(Options.class);
        Options opts = (Options) optParser.parse(args, true);

        final MemoryMXBean memoryBean = ManagementFactory.getMemoryMXBean();
		MemoryUsage memuse;

		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: initial: %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}
		try {
			Experiment.initialize(opts.config);
			Experiment experiment = Experiment.getInstance();
			System.out.println("experient loaded");
			System.out.flush();
			experiment.buildPrefixes();
			System.out.println("buildPrefixes() done");
			System.out.flush();
			experiment.buildWordPrefixes();
			System.out.println("buildWordPrefixes() done");
			System.out.flush();
			if (opts.meminfo) {
				memuse = memoryBean.getHeapMemoryUsage();
				System.out.printf("MEMUSE: %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
			}
			LanguageModel lm = experiment.getLM(opts.lm);
			lm.getHistoryTree();
			System.out.println("getLM() done");
			System.out.flush();
			/*
			if (lm.getBackoffLM() != null) {
				experiment.getLM(lm.getBackoffLM());
			}
			System.out.println("getLM(backoff) done");
			System.out.flush();
			*/
			//AbstractCostFunction costFunction = experiment.getCostFunction();
			
			experiment.closeXML();
			if (opts.meminfo) {
				memuse = memoryBean.getHeapMemoryUsage();
				System.out.printf("MEMUSE: %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
			}

			JobManager.initialize(opts.jobs);
			Thread thread = new Thread(JobManager.getInstance(), "Job Manager");
			thread.setDaemon(true);
			thread.start();
			
			File dbDir = new File(opts.dbdir);
			File dataDir = new File(opts.datadir);
			if (!dataDir.isDirectory()) {
				dataDir.mkdirs();
			}
			
			{
				
				Timer timer = new Timer(true);
				if (opts.meminfo) {
					TimerTask task = new TimerTask() {
						public void run() {
							MemoryUsage memuse = memoryBean.getHeapMemoryUsage();
							DateFormat df = DateFormat.getDateTimeInstance(DateFormat.SHORT, DateFormat.LONG);
							System.out.printf("MEMUSE: %s: %dM/%dM\n", df.format(new Date()), memuse.getUsed()/1048576, memuse.getMax()/1048576);					
						}
					};
					timer.scheduleAtFixedRate(task, 0, 10000);
				}
			
				TreeNgramDistribution.dumpNgramCounts = opts.dumpNgrams;
				
				NewEMSmoother smoother = new NewEMSmoother(lm, dbDir, dataDir, !opts.skipData, opts.dumpLambdasPrefix);
				if (!opts.skipData) {
					smoother.initialize(0);
				}
				
				if (opts.pruneTree) {
					smoother.pruneNodes();
				}
				
				smoother.doCrossValidatingSmoothing(opts.maxiter, opts.bucketSize, lm.getPruningThreshold(), 
						opts.useLambdas, (byte) opts.singleFold, opts.dummySmoother, opts.discountScale);
			}
			System.gc();
			if (opts.meminfo) {
				memuse = memoryBean.getHeapMemoryUsage();
				System.out.printf("MEMUSE: after gc %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);				
			}
			
			File backup = new File(lm.getHistoryTreeFilename() + ".bak");
			if (backup.exists()) {
				backup.delete();
			}
			new File(lm.getHistoryTreeFilename()).renameTo(backup);
			
			lm.saveHistoryTree();
			
			//experiment.saveConfig(opts.config);
			
			//smoother.saveSmoothedClusters();
			
		} catch(IOException e) {
			e.printStackTrace();
		}
	}

}
